#include "pch.h"
#include "FlvWriter.h"
#include "File.h"
#include "Utf8.h"


FlvWriter::FlvWriter(std::wstring outputPath, std::wstring sourcePath) {
    path = outputPath;
    auto dir = File::Directory(path);
    if (!File::Exist(dir)) {
        File::CreateDeepDirectory(dir);
    }
    if (File::Exist(outputPath)) {
        File::Delete(outputPath);
    }
    flv = std::ofstream(outputPath, std::fstream::binary);
    source = std::ifstream(sourcePath, std::fstream::binary);
}

FlvWriter::~FlvWriter() {
}

std::wstring FlvWriter::FullName() {
    return path;
}

std::ofstream& FlvWriter::BaseStream() {
    return flv;
}

void FlvWriter::WriteTag(void* tagPoint, int size) {
    flv.seekp(0, std::ios_base::end);
    flv.write( (const char*)tagPoint, size);
}

void FlvWriter::Parse() {
    meta =  Metadata();
    SeekLastTag();
    meta.Duration(GetTimeStamp() / double(1000));
    meta.LastTimesStamp(GetTimeStamp() / double(1000));
    SeekFirstTag();
    while (true) {
        if (GetTagType() != 9) {
            SeekNextTag();
            continue;
        }
        if (IsKeyFrame()) {
            auto info = ParseVideoTag();
            meta.VideoCodeCID(info[0]);
            break;
        }
    }
    ::SendMessage(hDlg, msg, 3, ProgressInterval() + Index());
    while (true) {
        if (GetTagType() != 8) {
            SeekNextTag();
            continue;
        }
        auto info = ParseAudioTag();
        meta.AudioCodeCID(info[0]) ;
        meta.AudioSampleRate(info[1]) ;
        meta.AudioSampleSize(info[2]);
        meta.Stereo(info[3] == 1 ? true : false);
        break;
    }
    ::SendMessage(hDlg, msg, 4, ProgressInterval() + Index());
    int i = 0;
    int samples[6] = {0};
    while (i < 6) {
        if (GetTagType() != 9 || GetTimeStamp() < 1) {
            SeekNextTag();
            continue;
        }
        samples[i] = GetTimeStamp();
        SeekNextTag();
        i++;
    }
    ::SendMessage(hDlg, msg, 5, ProgressInterval() + Index());
    meta.FrameRate(9000 / (samples[5] + samples[4] + samples[3] - samples[2] - samples[1] - samples[0])); //Framerate UNIT = fps
    const int breaking = 2000; //Breaking between two keyframes. UNIT = ms
    int lastTimestamp = -breaking;
    SeekFirstTag();
    while (IsNotEnd()) {
        if (GetTagType() != 9) {
            SeekNextTag();
            continue;
        }
        if (!IsKeyFrame()) {
            SeekNextTag();
            continue;
        }
        int timestamp = GetTimeStamp();
        if (timestamp - lastTimestamp >= breaking) {
            meta.keyframes.Add(Position(), GetTimeStamp() / double(1000));
            lastTimestamp = timestamp;
        }
        SeekNextTag();
    }
    ::SendMessage(hDlg, msg, 6, ProgressInterval() + Index());
    Position( meta.keyframes.filePositions.back());
    while (IsNotEnd()) {
        if (GetTagType() != 9) {
            SeekNextTag();
            continue;
        }
        if (IsKeyFrame()) {
            meta.LastKeyFrameTimeStamp(  GetTimeStamp() / double(1000));
        }
        SeekNextTag();
    }
    ::SendMessage(hDlg, msg, 7, ProgressInterval() + Index());
    meta.Width(0);
    meta.Height(0);
    meta.VideoDataRate( 0);
    meta.FileSize( 0);
    meta.AudioDelay( 0);
    meta.CanSeekToEnd(true);
    //meta.CreationDate(DateTime.Now.ToString("yyyy-MM-dd HH:mm:ss"));
    SYSTEMTIME time;
    GetLocalTime(&time);
    TCHAR date[128] = { 0 };
    _snwprintf(date, 128, L"%04d-%02d-%02d %02d:%02d:%02d", time.wYear, time.wMonth, time.wDay, time.wHour, time.wMinute, time.wSecond);
    meta.CreationDate(date);
    meta.MetaDataCreator(L"QSV2FLV v2.0 by Var");
    ::SendMessage(hDlg, msg, 8, ProgressInterval() + Index());
}

void FlvWriter::Output() {

    //FLV Base Info
    flv.seekp( 0, std::ios_base::beg);
    flv.write((char*)Metadata::Meta_1, sizeof(Metadata::Meta_1));
    flv.seekp(Metadata::DurationPosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.Duration()).data(), 8);
    //flv.Seek(Metadata.width_position, SeekOrigin.Begin);
    //flv.Write(DoubleToBytes(meta.width), 0, 8);
    //flv.Seek(Metadata.height_position, SeekOrigin.Begin);
    //flv.Write(DoubleToBytes(meta.height), 0, 8);
    flv.seekp(Metadata::VideoDataRatePosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.VideoDataRate()).data(), 8);

    flv.seekp(Metadata::FrameRatePosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.FrameRate()).data(), 8);

    flv.seekp( Metadata::VideoCodeCIDPosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.VideoCodeCID()).data(), 8);

    flv.seekp(Metadata::AudioSampleRatePosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.AudioSampleRate()).data(), 8);

    flv.seekp( Metadata::AudioSampleSizePosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.AudioSampleSize()).data(), 8);

    flv.seekp(Metadata::StereoPosition, std::ios_base::beg);
    flv.write((char*)BoolToBytes(meta.Stereo()).data(), 1);


    flv.seekp(Metadata::AudioCodeCIDPosition,  std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.AudioCodeCID()).data(), 8);

    //flv.Seek(Metadata.filesize_position, SeekOrigin.Begin);
    //flv.Write(DoubleToBytes(meta.filesize), 0, 8);
    flv.seekp(Metadata::LastTimeStampPosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.LastTimesStamp()).data(), 8);

    flv.seekp(Metadata::LastKeyFrameTimeStampPosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(meta.LastKeyFrameTimeStamp()).data(), 8);
    //flv.Seek(Metadata.audiodelay_position, SeekOrigin.Begin);
    //flv.Write(DoubleToBytes(meta.audiodelay), 0, 8);
    flv.seekp(Metadata::CanSeekToEndPosition, std::ios_base::beg);
    flv.write((char*)BoolToBytes(meta.CanSeekToEnd()).data(), 1);

    flv.seekp(Metadata::CreationDatePosition, std::ios_base::beg);
    flv.write((char*)StringToBytes(meta.CreationDate()).data(), 11);

    flv.seekp(Metadata::MetaDataCreatorPosition, std::ios_base::beg);
    flv.write((char*)ShorToBytes(meta.MetaDataCreator().size()).data(), 2);
    flv.write((char*)StringToBytes(meta.MetaDataCreator()).data(), meta.MetaDataCreator().length());

    //Keyframes
    flv.seekp(0, std::ios_base::end);
    flv.write((char*)Metadata::Meta_2, sizeof(Metadata::Meta_2));
    int headerSize = meta.keyframes.Count() * 18 + meta.MetaDataCreator().size() + 466;
    //filepositions
    flv.write((char*)IntToBytes(meta.keyframes.Count()).data(), 4);
    Byte zeroBuffer =  0 ;
    for (long i : meta.keyframes.filePositions) {
        flv.write(&zeroBuffer, 1);
        flv.write((char*)DoubleToBytes(i + headerSize).data(), 8);
    }
    //times
    flv.write((char*)Metadata::Meta_3, sizeof(Metadata::Meta_3));
    flv.write((char*)IntToBytes(meta.keyframes.Count()).data(), 4);
    for (long i : meta.keyframes.times) {
        flv.write(&zeroBuffer, 1);
        flv.write((char*)DoubleToBytes(i).data(), 8);
    }


    ::SendMessage(hDlg, msg, 9, ProgressInterval() + Index());
    //Pre-tag Size
    flv.seekp(Metadata::MetaDataSizePosition, std::ios_base::beg);
    flv.write((char*)IntToBytes(GetStreamLength(flv) - 18).data() + 1, 3);
    flv.seekp(0, std::ios_base::end);
    flv.write((char*)Metadata::Meta_4, sizeof(Metadata::Meta_4));
    flv.write((char*)IntToBytes(GetStreamLength(flv) - 13).data(), 4);

    //Copy Source Stream
    source.seekg(0, std::ios_base::beg);
    flv.seekp(0, std::ios_base::end);
    const int bufferSize = 10485760;//10MB
    // Byte buffer[bufferSize] = { 0 }; stack overflow error
    Byte* buffer = new Byte[bufferSize]();

    int size = GetStreamLength(source) % bufferSize;
    while (source.read((char*)buffer, bufferSize)) {
        flv.write((char*)buffer, bufferSize);
    }
    flv.write((char*)buffer, size);
    delete[] buffer;
    //Metadata.filesize
    flv.seekp(Metadata::FileSizePosition, std::ios_base::beg);
    flv.write((char*)DoubleToBytes(GetStreamLength(flv)).data(), 4);
    ::SendMessage(hDlg, msg, 10, ProgressInterval() + Index());

    flv.close();
    source.close();
}

std::streamsize FlvWriter::GetStreamLength(std::ifstream& fs) {
    std::ifstream::pos_type cur = fs.tellg();
    fs.seekg(0, std::ios_base::end);
    std::streamsize length = fs.tellg();
    fs.seekg(cur, std::ios_base::beg);
    return length;
}

std::streamsize FlvWriter::GetStreamLength(std::ofstream& fs) {
    std::ofstream::pos_type cur = fs.tellp();
    fs.seekp(0, std::ios_base::end);
    std::streamsize length = fs.tellp();
    fs.seekp(cur, std::ios_base::beg);
    return length;
}

//start
std::streampos FlvWriter::Position() {
    return source.tellg();
}

void FlvWriter::Position(std::streampos position) {
    source.seekg(position);
}

void FlvWriter::SeekFirstTag() {
    source.seekg(0, std::ios_base::beg);
}

void FlvWriter::SeekNextTag() {
    Byte buffer[4] = { 0 };
    source.read((char*)buffer, 4);
    if (!(buffer[0] == 8 || buffer[0] == 9)) {
        throw std::exception("Unknow tag.");
    }
    int dataSize = (0xFF & buffer[1]) << 16 | (0xFF & buffer[2]) << 8 | (0xFF & buffer[3]);
    try {
        source.seekg(dataSize + 11, std::ios_base::cur);
    } catch (std::exception e) {
        throw;
    }
}

void FlvWriter::SeekLastTag() {
    Byte buffer[4] = { 0 };
    source.seekg(-4, std::ios_base::end);
    source.read((char*)buffer, 4);
    int tagSize = (0xFF & buffer[0]) << 24 | (0xFF & buffer[1]) << 16 | (0xFF & buffer[2]) << 8 | (0xFF & buffer[3]);
    source.seekg(-tagSize - 4, std::ios_base::end);
}

bool FlvWriter::IsNotEnd() {
    //std::ifstream::pos_type cur = source.tellg();
    //source.seekg(0, std::ios_base::beg);
    //bool length = source.rdbuf()->in_avail();
    //source.seekg(cur, std::ios_base::beg);
    return source.tellg() < GetStreamLength(source);
}

std::vector<int> FlvWriter::ParseAudioTag() {
    Byte buffer =  0 ;
    source.read((char*)&buffer, 1);
    int num = *reinterpret_cast<unsigned char*>(&buffer);
    if (num != 8) {
        throw std::exception("Not a audio tag.");
    }
    source.seekg(10, std::ios_base::cur);
    source.read((char*)&buffer, 1);
    num = *reinterpret_cast<unsigned char*>(&buffer);
    source.seekg(-12, std::ios_base::cur);
    int SoundFormat = num >> 4;
    int SoundRate = num >> 2 & 3;
    int SoundSize = num >> 1 & 1;
    int SoundType = num & 1;
    return { SoundFormat, SoundRate, SoundSize, SoundSize };
}

std::vector<int> FlvWriter::ParseVideoTag() {
    Byte buffer =  0 ;
    source.read((char*)&buffer, 1);
    int num = *reinterpret_cast<unsigned char*>(&buffer);
    if (num != 9) {
        throw std::exception("Not a video tag.");
    }
    source.seekg(10, std::ios_base::cur);
    source.read((char*)&buffer, 1);
    num = *reinterpret_cast<unsigned char*>(&buffer);
    source.seekg(-12, std::ios_base::cur);
    if ((num & 17) != 17) {
        throw std::exception("Not a keyframe.");
    }
    int CodecID = num & 15;
    return { CodecID };
}

bool FlvWriter::IsKeyFrame() {
    Byte buffer = 0;
    source.read((char*)&buffer, 1);
    int num = *reinterpret_cast<unsigned char*>(&buffer);
    if (num != 9) {
        throw std::exception("Not a video tag.");
    }
    source.seekg(10, std::ios_base::cur);
    source.read((char*)&buffer, 1);
    num = *reinterpret_cast<unsigned char*>(&buffer);
    source.seekg(-12, std::ios_base::cur);
    return (num & 17) == 17 ? true : false;
}

int FlvWriter::GetTimeStamp() {
    Byte buffer[4] = { 0 };
    source.seekg( 4, std::ios_base::cur);
    source.read((char*)buffer, 4);
    source.seekg(-8, std::ios_base::cur);
    int timestamp =
        (0xFF & buffer[0]) << 16 | (0xFF & buffer[1]) << 8 | (0xFF & buffer[2]) | (0xFF & buffer[3]) << 24;
    return timestamp;
}

int FlvWriter::GetTagType() {
    Byte buffer = 0;
    source.read((char*)&buffer, 1);
    source.seekg(-1, std::ios_base::cur);
    int num = *reinterpret_cast<unsigned char*>(&buffer);
    return num;
}

std::vector<unsigned char> FlvWriter::IntToBytes(int num) {
    std::vector<unsigned char> bytes;
    std::copy(static_cast<const unsigned char*>(static_cast<const void*>(&num)),
              static_cast<const unsigned char*>(static_cast<const void*>(&num)) + sizeof(num),
              std::back_inserter(bytes));
    std::reverse(std::begin(bytes), std::end(bytes));
    return bytes;
}

std::vector<unsigned char> FlvWriter::ShorToBytes(int num) {
    std::vector<unsigned char> bytes;
    std::copy(static_cast<const unsigned char*>(static_cast<const void*>(&num)),
              static_cast<const unsigned char*>(static_cast<const void*>(&num)) + sizeof(num),
              std::back_inserter(bytes));
    auto temp = bytes[0];
    bytes[0] = bytes[1];
    bytes[1] = temp;
    return bytes;
}

std::vector<unsigned char> FlvWriter::DoubleToBytes(double num) {
    std::vector<unsigned char> bytes;
    std::copy(static_cast<const unsigned char*>(static_cast<const void*>(&num)),
              static_cast<const unsigned char*>(static_cast<const void*>(&num)) + sizeof(num),
              std::back_inserter(bytes));
    std::reverse(std::begin(bytes), std::end(bytes));
    return bytes;
}

std::vector<unsigned char> FlvWriter::StringToBytes(std::wstring str) {
    std::vector<unsigned char> bytes;
    Utf8::Utf16ToUtf8(str);
    auto ansiStr = Utf8::Utf16ToUtf8(str);
    bytes.resize(ansiStr.size());
    std::copy(std::begin(ansiStr), std::end(ansiStr), std::begin(bytes));
    return bytes;
}

std::vector< unsigned char> FlvWriter::BoolToBytes(bool bl) {
    std::vector<unsigned char> bytes;
    std::copy(static_cast<const unsigned char*>(static_cast<const void*>(&bl)),
              static_cast<const unsigned char*>(static_cast<const void*>(&bl)) + sizeof(bl),
              std::back_inserter(bytes));
    return bytes;
}
